from flask_wtf import FlaskForm
from wtforms import StringField, PasswordField, BooleanField, SelectField, SubmitField, TextAreaField, DateTimeField, FloatField
from wtforms.validators import DataRequired, Email, Length, EqualTo, ValidationError, Optional
from .models.user import User

class RegistrationForm(FlaskForm):
    username = StringField('用户名', validators=[
        DataRequired(message='用户名不能为空'),
        Length(min=3, max=20, message='用户名长度必须在3-20个字符之间')
    ])
    email = StringField('邮箱', validators=[
        DataRequired(message='邮箱不能为空'),
        Email(message='请输入有效的邮箱地址')
    ])
    password = PasswordField('密码', validators=[
        DataRequired(message='密码不能为空'),
        Length(min=6, message='密码长度不能少于6个字符')
    ])
    confirm_password = PasswordField('确认密码', validators=[
        DataRequired(message='请确认密码'),
        EqualTo('password', message='两次输入的密码不一致')
    ])
    role = SelectField('角色', choices=[
        ('researcher', '研究人员'),
        ('vet', '兽医'),
        ('admin', '管理员')
    ], validators=[DataRequired(message='请选择用户角色')])
    submit = SubmitField('注册')

    def validate_username(self, field):
        if User.query.filter_by(username=field.data).first():
            raise ValidationError('该用户名已被使用')

    def validate_email(self, field):
        if User.query.filter_by(email=field.data).first():
            raise ValidationError('该邮箱已被注册')

class LoginForm(FlaskForm):
    username = StringField('用户名', validators=[DataRequired(message='请输入用户名')])
    password = PasswordField('密码', validators=[DataRequired(message='请输入密码')])
    remember_me = BooleanField('记住我')
    submit = SubmitField('登录')

class ForgotPasswordForm(FlaskForm):
    email = StringField('邮箱', validators=[
        DataRequired(message='请输入邮箱'),
        Email(message='请输入有效的邮箱地址')
    ])
    submit = SubmitField('发送重置链接')

class ResetPasswordForm(FlaskForm):
    password = PasswordField('新密码', validators=[
        DataRequired(message='请输入新密码'),
        Length(min=6, message='密码长度不能少于6个字符')
    ])
    confirm_password = PasswordField('确认密码', validators=[
        DataRequired(message='请确认密码'),
        EqualTo('password', message='两次输入的密码不一致')
    ])
    submit = SubmitField('重置密码')

class ProfileForm(FlaskForm):
    name = StringField('真实姓名')
    phone = StringField('电话号码')
    organization = StringField('所属组织')
    email = StringField('邮箱', validators=[Email(message='请输入有效的邮箱地址')])
    submit = SubmitField('更新资料')

class PathogenForm(FlaskForm):
    name = StringField('病原体名称', validators=[DataRequired(message='请输入病原体名称')])
    type = SelectField('类型', choices=[
        ('virus', '病毒'),
        ('bacteria', '细菌'),
        ('parasite', '寄生虫'),
        ('fungus', '真菌')
    ], validators=[DataRequired(message='请选择病原体类型')])
    isolation_date = DateTimeField('分离时间', format='%Y-%m-%d', 
        validators=[DataRequired(message='请选择分离时间')])
    isolation_location = StringField('分离地点', 
        validators=[DataRequired(message='请输入分离地点')])
    latitude = FloatField('纬度', validators=[Optional()])
    longitude = FloatField('经度', validators=[Optional()])
    gene_sequence = TextAreaField('基因序列')
    submit = SubmitField('保存')

class PathogenSearchForm(FlaskForm):
    keyword = StringField('关键词')
    type = SelectField('类型', choices=[
        ('', '所有类型'),
        ('virus', '病毒'),
        ('bacteria', '细菌'),
        ('parasite', '寄生虫'),
        ('fungus', '真菌')
    ], default='')
    date_from = DateTimeField('开始日期', format='%Y-%m-%d', validators=[Optional()])
    date_to = DateTimeField('结束日期', format='%Y-%m-%d', validators=[Optional()])
    submit = SubmitField('搜索')

class UserForm(FlaskForm):
    username = StringField('用户名', validators=[DataRequired(), Length(min=3, max=20)])
    email = StringField('邮箱', validators=[DataRequired(), Email()])
    password = PasswordField('密码', validators=[Optional(), Length(min=6)])
    is_admin = BooleanField('管理员权限')
    submit = SubmitField('提交') 